package com.aqie.arithmetic.datastructure;

/**
 * 红黑树
 */
public class RedBlackBST<Key extends Comparable<Key>, Value> {
    private Node root;
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private class Node{
        Key key;
        Value value;
        Node left, right;
        // 子树中结点总数
        int N;
        // 父节点指向它的链接的颜色
        boolean color;

        Node(Key key, Value value, int N, boolean color){
            this.color = color;
            this.key = key;
            this.value = value;
            this.N = N;
        }
        public Node rotateLeft(Node h){
            Node x = h.right;
            h.right = x.left;
            x.left = h;
            x.color = h.color;
            h.color = RED;
            x.N = h.N;
            h.N = 1 + size(h.left) + size(h.right);
            return x;
        }

        public Node rotateRight(Node h){
            Node x = h.left;
            h.left = x.right;
            x.right = h;
            x.color = h.color;
            h.color = RED;
            x.N = h.N;
            h.N = 1 + size(h.left) + size(h.right);
            return x;
        }
    }

    private boolean isRed(Node x){
        if (x == null) return false;
        return x.color == RED;
    }

    public int size(){
        return size(root);
    }

    private int size(Node x){
        if (x == null) return 0;
        else return x.N;
    }

    public boolean isEmpty(){
        return size() == 0;
    }


    private void flipColors(Node h){
        h.color = RED;
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

    public void put(Key key, Value value){
        root = put(root, key, value);
        root.color = BLACK;
    }

    private Node put(Node h, Key key, Value value) {
        if (h == null){
            return new Node(key, value, 1, RED);
        }
        int cmp = key.compareTo(h.key);
        if (cmp < 0) h.left = put(h.left, key, value);
        else if (cmp > 0) h.right = put(h.right, key, value);
        else h.value = value;

        if (isRed(h.right) && !isRed(h.left)) h.rotateLeft(h);
        if (isRed(h.left) && isRed(h.left.left)) h.rotateRight(h);
        if (isRed(h.left) && isRed(h.right)) flipColors(h);
        h.N = size(h.left) + size(h.right) + 1;
        return h;
    }

    // 删除最小键
    private Node moveRedLeft(Node h){
        // 假设节点 h 为红色，h.left 和 h.left.left 都是黑色
        // 将 h.left 或 h.left的子节点之一变红
        flipColors(h);
        if (isRed(h.right.left)){
            h.right = h.rotateRight(h.right);
            h = h.rotateLeft(h);
        }
        return h;
    }

    public void deleteMin(){
        if (!isRed(root.left) && !isRed(root.right)){
            root.color = RED;
        }
        root = deleteMin(root);
        if (!isEmpty()){
            root.color = BLACK;
        }
    }

    private Node deleteMin(Node h){
        if (h.left == null){
            return null;
        }
        if (!isRed(h.left) && !isRed(h.left.left)){
            h = moveRedLeft(h);
        }
        h.left = deleteMin(h.left);
        return balance(h);
    }

    private Node balance(Node h){
        if (isRed(h.right)) h = h.rotateLeft(h);
        if (isRed(h.right) && !isRed(h.left)) h.rotateLeft(h);
        if (isRed(h.left) && isRed(h.left.left)) h.rotateRight(h);
        if (isRed(h.left) && isRed(h.right)) flipColors(h);
        h.N = size(h.left) + size(h.right) + 1;
        return h;
    }

    // 删除最大键
    private Node moveRedRight(Node h){
        // 假设结点h为红色， h.right 和 h.right.left都是黑色
        // 将h.right 或h.right的子节点之一变红
        flipColors(h);
        if (!isRed(h.left.left)){
            h = h.rotateRight(h);
        }
        return h;
    }

    public void deleteMax(){
        if (!isRed(root.left) && !isRed(root.right)){
            root.color = RED;
        }
        root = deleteMax(root);
        if (!isEmpty()) root.color = BLACK;
    }

    private Node deleteMax(Node h){
        if (isRed(h.left)){
            h = h.rotateRight(h);
        }
        if (h.right == null){
            return null;
        }
        if (!isRed(h.right) && !isRed(h.right.left)){
            h = moveRedRight(h);
        }
        h.right = deleteMax(h.right);
        return balance(h);
    }

    // 删除操作
    public void delete(Key key){
        if (!isRed(root.left) && !isRed(root.right)){
            root.color = RED;
        }
        root = delete(root, key);
        if (!isEmpty()) root.color = BLACK;
    }

    private Node delete(Node h, Key key) {
        if (key.compareTo(h.key) < 0){
            if (!isRed(h.left) && !isRed(h.left.left)){
                h = moveRedLeft(h);
            }
            h.left = delete(h.left, key);
        }else {
            if (isRed(h.left)){
                h = h.rotateRight(h);
            }

            if (key.compareTo(h.key) == 0 && (h.right == null)){
                return null;
            }

            if (!isRed(h.right) && !isRed(h.right.left)){
                h = moveRedRight(h);
            }
            if (key.compareTo(h.key) == 0){
                h.value = get(h.right, min(h.right).key);
                h.key = min(h.right).key;
                h.right = deleteMin(h.right);
            }else{
                h.right = delete(h.right, key);
            }
        }
        return balance(h);
    }


    public Value get(Key key){
        // 3.3.1
        return get(root, key);
    }

    private Value get(Node x, Key key) {
        if (x == null) return null;
        int cmp = key.compareTo(x.key);
        if (cmp < 0) return get(x.left, key);
        else if (cmp > 0) return get(x.right, key);
        else return x.value;
    }

    public Key min(){
        return min(root).key;
    }

    private Node min(Node x) {
        if (x.left == null) return x;
        return min(x.left);
    }
}
